import numpy as np
import cv2 as cv

color = np.random.randint(0, 255, (100, 3))

# MOG背景分割器
mog = cv.createBackgroundSubtractorMOG2(detectShadows=True)

cap = cv.VideoCapture('./vtest.avi')
ret, frame = cap.read()


def labelTargets(src, mask, thresh=100):
    seg = np.copy(mask)
    # 寻找轮廓
    _, cnts, hier = cv.findContours(seg, cv.RETR_EXTERNAL,
                                    cv.CHAIN_APPROX_SIMPLE)
    count = 0
    for i in range(len(cnts)):
        c = cnts[i]
        area = cv.contourArea(c)
        if area < thresh:  # 面积小于阈值的当噪声处理
            continue
        count += 1  # 符合条件的轮廓计数加一
        x, y, w, h = cv.boundingRect(c)  # 寻找矩形的长宽高
        cv.rectangle(src, (x, y), (x + w, y + h),
                     color[np.random.randint(100)].tolist(), 2)  # 画矩形
        print('目标位置: ' + '(' + str(x) + ',' + str(y) + ')')
        print('目标面积：', area)
        # 在矩形左上角标记个数
        cv.putText(src, str(count), (x, y), cv.FONT_HERSHEY_PLAIN, 0.5,
                   (0, 255, 0))
    return count


while ret:
    # 将原图缩小一半
    image = cv.resize(frame, (int(frame.shape[0] / 2), int(frame.shape[1] / 2)),
                      interpolation=cv.INTER_LINEAR)
    # 通过MOG背景分割器得到前景滤镜
    fgmask = mog.apply(image)
    # 把前景滤镜二值化
    fgmask = cv.threshold(np.copy(fgmask), 30, 255, cv.THRESH_BINARY)[1]
    # 创建一个空的前景图
    foreGround = np.zeros(image.shape, image.dtype)
    # 把前景滤镜的维度从灰度图转为和彩色图一样
    mymask = cv.cvtColor(fgmask, cv.COLOR_GRAY2BGR)
    # 把原图拷贝到前景图，除了滤镜部分
    np.copyto(foreGround, image, where=mymask.astype(bool))
    # 把找到的目标标记出来
    cnt = labelTargets(image, fgmask)
    # 得到背景图像
    back = mog.getBackgroundImage()
    print('检测到的目标数:', cnt)

    cv.imshow('origin', image)
    cv.imshow('forgemask', fgmask)
    cv.imshow('background', back)
    cv.imshow("forgeground", foreGround)
    ret, frame = cap.read()  # 读取视频帧数据
    if cv.waitKey(100) & 0xff == ord("q"):
        break

cap.release()
cv.destroyAllWindows()
